from diagrams import Diagram, Cluster
from diagrams.generic.network import Switch
from diagrams.generic.storage import Storage
from diagrams import Diagram, Cluster
from diagrams.aws.compute import EC2
from diagrams.aws.general import TraditionalServer
import boto3
STATUS_COLORS = {
    "online": "lightgreen",
    "offline": "gray",
    "warning": "orange",
    "error": "red"
}



# 获取EC2实例状态

from diagrams import Diagram, Cluster
from diagrams.generic.network import Switch
from diagrams.generic.storage import Storage
from diagrams.aws.compute import EC2


VERTICAL_CONFIG = {
    "direction": "LR",  # 从左到右的垂直布局
    "graph_attr": {"rankdir": "BT"}  # 节点从上到下排列
}


STATUS_COLORS = {
    "online": "#4CAF50",  # 运行中
    "offline": "#9E9E9E", # 离线
    "warning": "#FFC107", # 警告
    "error": "#F44336",   # 故障
    "maintenance": "#2196F3" # 维护中
}

def create_device(name, status, u_height=1, node_type=EC2):
    return node_type(f"{name}\n{u_height}U", **{
        "fillcolor": STATUS_COLORS[status],
        "style": "filled",
        "fontcolor": "white" if status in ["online","error"] else "black"
    })

with Diagram("rack", show=False,direction="TB"):
    with Cluster("Rack A (42U)"):
        # 网络设备层 (顶部4U)
        tor_switch = Switch("TOR Switch\n10Gbps\n4U")

        # 服务器层 (中间34U)
        servers = [
            create_device("Web-01", "online", 2),
            create_device("DB-01", "warning", 4),
            create_device("Cache-01", "error", 2),
            create_device("Backup-01", "offline", 4)
        ]

        # 存储层 (底部4U)
        storage = Storage("NAS-01\nRAID6\n4U", **{
            "fillcolor": "#607D8B",
            "style": "filled"
        })

        # 连接关系
       # tor_switch >> servers
      #  servers[1] >> storage
      #  servers[3] - storage  # 备用连接

